/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import static org.sss.common.Utils.loadMapperClass;
import static org.sss.common.Utils.loadReducerClass;

import java.io.IOException;
import java.rmi.RemoteException;
import java.rmi.registry.Registry;
import java.rmi.server.UnicastRemoteObject;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import org.slf4j.Logger;
import org.sss.common.CommonStatus;
import org.sss.common.GetStatus;
import org.sss.common.Inlet;
import org.sss.common.JobDoneListener;
import org.sss.common.Outlet;
import org.sss.common.Reflection;
import org.sss.common.SssServerRemote;
import org.sss.common.TaskInfo;
import org.sss.common.TaskStatus;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;

public class SssServerImpl extends UnicastRemoteObject implements SssServerRemote {
  private static final long serialVersionUID = 1L;
  private static Logger logger = SssServer.getLogger(SssServerImpl.class);

  private final Map<UUID, TaskSet> taskSetMap = new HashMap<UUID, TaskSet>();
  private final Map<UUID, TaskStatus> taskStatusMap = new HashMap<UUID, TaskStatus>();
  private Registry registry;
  private SssServerResource resource;
  private void init(int id) {
    logger.debug("init id {}", id);
    try {
      this.resource = new SssServerResource(id);
    }
    catch (Exception e) {
      logger.error("failed in initializing server resources, exit.", e);
      System.exit(-1);
    }
  }
  public SssServerImpl() throws Exception {
    super();
    SssServer.initLogging();
    init(org.sss.util.Misc.randomInt());
    logger.info("this server is created by default constructor, not registered on rmi.");
  }
  public SssServerImpl(int id) throws Exception {
    super();
    SssServer.initLogging();
    init(id);
    logger.info("this server is created with id ({}), not registered on rmi.");
  }
  public SssServerImpl(Registry registry, int id) throws RemoteException {
    super();
    this.registry = registry;
    init(id);
  }
  public SssServerImpl(Registry registry, SssServerResource res) throws RemoteException {
    super();
    this.registry = registry;
    this.resource = res;
  }
  public SssServerImpl(SssServerResource res) throws Exception {
    super();
    SssServer.initLogging();
    this.resource = res;
    logger.info("this server is created with SssServerResource object, not registered on rmi.");
  }
  @Override
  public void startTaskSet(UUID taskSetID) throws RemoteException {
    synchronized (taskSetMap) {
      if (taskSetMap.containsKey(taskSetID)) {
        throw new RemoteException("Task set ID \"" + taskSetID + "\" has already been used.");
      }
      taskSetMap.put(taskSetID, new TaskSet());
    }
  }
  @Override
  public void endTaskSet(UUID taskSetID) throws RemoteException {
    synchronized (taskSetMap) {
      synchronized (taskStatusMap) {
        TaskSet taskSet = getTaskSet(taskSetID);
        for (UUID jobID: taskSet.getJobIDList()) {
          TaskStatus status = taskStatusMap.get(jobID);
          assert status != null: "TaskStatus object disappeared";
          status.stop();
          taskStatusMap.remove(jobID);
          cancel(jobID);
        }
        taskSetMap.remove(taskSetID);
      }
    }
  }
  @Override
  public void appendJarFile(UUID taskSetID, byte[] fileBuffer) throws RemoteException {
    synchronized (taskSetMap) {
      TaskSet ts = getTaskSet(taskSetID);
      try {
        ts.getClassLoader().appendJar(fileBuffer);
      } catch (IOException e) {
        throw new RemoteException("", e);
      }
    }
  }
  @Override
  public TaskStatus getTaskStatus(UUID jobID) throws RemoteException {
    logger.info("getTaskStatus {}", jobID);
    synchronized (taskStatusMap) {
      TaskStatus s = taskStatusMap.get(jobID);
      if (s != null) {
        return s.getCopy();
      }
      else {
        return null;
      }
    }
  }
  @Override
  public TaskInfo getTaskInfo() throws RemoteException {
    TaskInfo ti = new TaskInfo();
    synchronized (taskSetMap) {
      for (Map.Entry<UUID, TaskSet> e: taskSetMap.entrySet()) {
        ti.taskMap.put(e.getKey(), e.getValue().getJobIDList());
      }
    }
    return ti;
  }
  @Override
  public void submitMapTask(UUID taskSetID,
                            UUID jobID,
                            StorageNode storageNode,
                            Configuration conf,
                            List<Inlet> inletList,
                            String mapperClassName,
                            List<Outlet> outletList,
                            JobDoneListener<CommonStatus> doneListener)
    throws RemoteException {
    printTaskInfoToLog("submitMapTask", Reflection.ProcType.MAP,
        taskSetID, storageNode, inletList, mapperClassName, outletList);
    TaskSet taskSet = getTaskSet(taskSetID);
    conf.setClassLoader(taskSet.getClassLoader());
    taskSet.addJobID(jobID);
    TaskStatus taskStatus;
    synchronized (taskStatusMap) {
      taskStatus = taskStatusMap.get(jobID);
      if (taskStatus == null) {
        taskStatus = new TaskStatus(jobID);
        taskStatusMap.put(jobID, taskStatus);
      }
    }
    try {
      if (inletList.size() != 1) {
        throw new SssException("unsupported inletList size");
      }
      MapTask task = new MapTask(taskSet,
          jobID, resource, conf, storageNode, inletList.get(0),
          loadMapperClass(taskSet.getClassLoader(), mapperClassName),
          outletList, JobDoneCallback.wrap(doneListener, taskStatus));
      task.submit();
    }
    catch (SssException e) {
      throw new RemoteException("", e);
    }
  }
  @Override
  public void submitReduceTask(UUID taskSetID,
                               UUID jobID,
                               StorageNode storageNode,
                               Configuration conf,
                               List<Inlet> inletList,
                               String reducerClassName,
                               List<Outlet> outletList,
                               JobDoneListener<CommonStatus> doneListener)
    throws RemoteException {
    printTaskInfoToLog("submitReduceTask", Reflection.ProcType.REDUCE,
        taskSetID, storageNode, inletList, reducerClassName, outletList);
    TaskSet taskSet = getTaskSet(taskSetID);
    conf.setClassLoader(taskSet.getClassLoader());
    taskSet.addJobID(jobID);
    TaskStatus taskStatus;
    synchronized (taskStatusMap) {
      taskStatus = taskStatusMap.get(jobID);
      if (taskStatus == null) {
        taskStatus = new TaskStatus(jobID);
        taskStatusMap.put(jobID, taskStatus);
      }
    }
    try {
      ReduceTask task = new ReduceTask(taskSet, jobID,
          resource, conf, storageNode, inletList.get(0),
          loadReducerClass(taskSet.getClassLoader(), reducerClassName),
          outletList, JobDoneCallback.wrap(doneListener, taskStatus));
      task.submit();
    }
    catch (SssException e) {
      throw new RemoteException("", e);
    } finally {
      logger.debug("submitReduceTask end.");
    }
  }
  @Override
  public void submitMMapTask(UUID taskSetID, UUID jobID,
    StorageNode storageNode, Configuration conf, List<Inlet> inletList,
    String mapperClassName, List<Outlet> outletList,
    JobDoneListener<CommonStatus> doneListener) throws RemoteException {
    printTaskInfoToLog("submitMMapTask", Reflection.ProcType.MAP,
        taskSetID, storageNode, inletList, mapperClassName, outletList);

    TaskSet taskSet = getTaskSet(taskSetID);
    conf.setClassLoader(taskSet.getClassLoader());
    taskSet.addJobID(jobID);
    TaskStatus taskStatus;
    synchronized (taskStatusMap) {
      taskStatus = taskStatusMap.get(jobID);
      if (taskStatus == null) {
        taskStatus = new TaskStatus(jobID);
        taskStatusMap.put(jobID, taskStatus);
      }
    }
    try {
      if (inletList.size() != 2) {
        throw new SssException("unsupported inletList size");
      }
      MergeMapTask task = new MergeMapTask(taskSet,
          jobID, resource, conf, storageNode, inletList.get(0), inletList.get(1),
          loadMapperClass(taskSet.getClassLoader(), mapperClassName),
          outletList,
          JobDoneCallback.wrap(doneListener, taskStatus));
      task.submit();
    }
    catch (SssException e) {
      throw new RemoteException("", e);
    }
    finally {
      logger.debug("submitMMapTask end.");
    }
  }
  @Override
  public void submitMReduceTask(UUID taskSetID, UUID jobID,
    StorageNode storageNode, Configuration conf, List<Inlet> inletList,
    String reducerClassName, List<Outlet> outletList,
    JobDoneListener<CommonStatus> doneListener) throws RemoteException {
    printTaskInfoToLog("submitMReduceTask", Reflection.ProcType.REDUCE,
        taskSetID, storageNode, inletList, reducerClassName, outletList);

    TaskSet taskSet = getTaskSet(taskSetID);
    conf.setClassLoader(taskSet.getClassLoader());
    taskSet.addJobID(jobID);
    TaskStatus taskStatus;
    synchronized (taskStatusMap) {
      taskStatus = taskStatusMap.get(jobID);
      if (taskStatus == null) {
        taskStatus = new TaskStatus(jobID);
        taskStatusMap.put(jobID, taskStatus);
      }
    }
    try {
      if (inletList.size() != 2) {
        throw new SssException("unsupported inletList size");
      }
      MergeReduceTask task = new MergeReduceTask(taskSet,
        jobID, resource, conf, storageNode, inletList.get(0), inletList.get(1),
        loadReducerClass(taskSet.getClassLoader(), reducerClassName),
        outletList, JobDoneCallback.wrap(doneListener, taskStatus));
      task.submit();
    }
    catch (SssException e) {
      throw new RemoteException("", e);
    }
    finally {
      logger.debug("submitMReduceTask end.");
    }
  }
  @Override
  public void submitCountTask(UUID jobID,
      StorageNode storageNode,
      Configuration conf, GroupID targetGroupID,
      JobDoneListener<GetStatus> doneListener) throws RemoteException {
    logger.info("submitCountTask start.");
    logger.info("  storageNode {}:{}", storageNode.getHostName(), storageNode.getPort());
    logger.info("  inputGroup {}", targetGroupID);

    try {
      CountTask task = new CountTask(
          jobID, resource, conf, storageNode, targetGroupID,
          JobDoneCallback.wrap(doneListener));
      task.submit();
    }
    catch (SssException e) {
      throw new RemoteException("", e);
    }
    finally {
      logger.debug("submitCountTask end.");
    }
  }

  @Override
  public Set<GroupID> getAllTuleGroupIDs() throws RemoteException {
    try {
      return CommandGetAllTupleGroups.execute(
          resource.getIOTPRead(),
          resource.getStorageNodeList());
    }
    catch (Exception e) {
      throw new RemoteException("Can't get tuple group ID list.", e);
    }
  }

  @Override
  public long removeTupleGroup(GroupID gid) throws RemoteException {
    try {
      return CommandRemoveTupleGroup.execute(resource.getIOTPWrite(),
                                resource.getStorageNodeList(), gid);
    }
    catch (Exception e) {
      throw new RemoteException("Rangeout failed.", e);
    }
  }

  @Override
  public void shutdown() throws RemoteException {
    logger.info("received shutdown request");
    try {
      if (registry != null) {
        registry.unbind("SssServerImpl");
        UnicastRemoteObject.unexportObject(this, true);
      }
    }
    catch (Exception e) {
      logger.error("failed in unbinding `SssServerImpl'");
      logger.error("{}", e);
    }
    resource.shutdown();
  }

  @Override
  public void cancel(UUID id) throws RemoteException {
    List<Task> ts = Task.get(id);

    if (ts.size() == 0) {
      logger.info("There is no task whose id is {}. Maybe the task already has completed.", id);
    } else {
      logger.info("Task({}) is canceled.", id);
      for (Task t: ts) {
        t.cancel();
      }
    }
  }

  @Override
  public void aliveCheck() throws RemoteException {
    // Do nothig.
  }

  private TaskSet getTaskSet(UUID id) throws RemoteException {
    synchronized (taskSetMap) {
      TaskSet ts = taskSetMap.get(id);
      if (ts == null) {
        throw new RemoteException("Unknown task set " + id);
      }
      return ts;
    }
  }

  private static void printTaskInfoToLog(String name,
                                         Reflection.ProcType type,
                                         UUID taskSetID,
                                         StorageNode storageNode,
                                         List<Inlet> inletList,
                                         String userClass,
                                         List<Outlet> outletList) {
    logger.info("{} start.", name);
    logger.info("  task set {}", taskSetID);
    logger.info("  storageNode {}:{}", storageNode.getHostName(), storageNode.getPort());
    for (Inlet inlet: inletList) {
      logger.info("  inputGroupID {}", inlet.GID);
      logger.info("  inputKeyClassName {}", inlet.keyClassName);
      logger.info("  inputValueClassName {}", inlet.valueClassName);
    }
    switch (type) {
    case MAP:    logger.info("  mapperClassName {}",  userClass); break;
    case REDUCE: logger.info("  reducerClassName {}", userClass); break;
    }
    for (Outlet outlet: outletList) {
      logger.info("  outputGroupID {}", outlet.GID);
      logger.info("  partitionerClassName {}", outlet.partitionerClassName);
      logger.info("  combinerClassName {}", outlet.combinerClassName);
    }
  }
}
